import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import math
import random
import os
import time

class PSO:
    #初始化
    def __init__(self, D, N, M, p_low, p_up, v_low, v_high, w = 1., c1 = 0.5, c2 = 0.5):
        self.w = w  # 惯性权值
        self.c1 = c1  # 个体学习因子
        self.c2 = c2  # 群体学习因子
        self.D = D  # 粒子维度
        self.N = N  # 粒子群规模，初始化种群个数
        self.M = M  # 最大迭代次数
        self.p_range = [p_low, p_up]  # 粒子位置的约束范围
        self.v_range = [v_low, v_high]  # 粒子速度的约束范围
        self.x = np.zeros((self.N, self.D))  # 所有粒子的位置
        self.v = np.zeros((self.N, self.D))  # 所有粒子的速度
        self.p_best = np.zeros((self.N, self.D))  # 每个粒子的最优位置
        self.g_best = np.zeros((1, self.D))[0]  # 种群（全局）的最优位置
        self.p_bestFit = np.zeros(self.N)  # 每个粒子的最优适应值
        self.g_bestFit = float('Inf')  # float('-Inf')，始化种群（全局）的最优适应值，由于求极小值，故初始值给大，向下收敛，这里默认优化问题中只有一个全局最优解

        # 初始化所有个体和全局信息
        for i in range(self.N):
            for j in range(self.D):
                self.x[i][j] = random.uniform(self.p_range[0][j], self.p_range[1][j])
                self.v[i][j] = random.uniform(self.v_range[0], self.v_range[1])
            self.p_best[i] = self.x[i]  # 保存个体历史最优位置，初始默认第0代为最优
            fit = self.fitness(self.p_best[i])
            self.p_bestFit[i] = fit  # 保存个体历史最优适应值
            if fit < self.g_bestFit:  # 寻找并保存全局最优位置和适应值
                self.g_best = self.p_best[i]
                self.g_bestFit = fit
    #适应性函数
    def fitness(self,x):
        o = np.sum(x * x)
        return o
    #迭代函数
    def update(self,ax1):
        x=[]
        y=[]
        for i in range(self.N):
            self.v[i] = self.w * self.v[i] + self.c1 * random.uniform(0, 1) * (self.p_best[i] - self.x[i]) + self.c2 * random.uniform(0, 1) * (self.g_best - self.x[i])
            #速度限制
            for j in range(self.D):
                if self.v[i][j] < self.v_range[0]:
                    self.v[i][j] = self.v_range[0]
                if self.v[i][j] > self.v_range[1]:
                    self.v[i][j] = self.v_range[1]
            #更新位置
            self.x[i] = self.x[i] + self.v[i]
            #位置限制
            for j in range(self.D):
                if self.x[i][j] < self.p_range[0][j]:
                    self.x[i][j] = self.p_range[0][j]
                if self.x[i][j] > self.p_range[1][j]:
                    self.x[i][j] = self.p_range[1][j]
            #更新个体和全局历史最优位置及适应值
            _fit = self.fitness(self.x[i])
            x.append(self.x[i][0])
            y.append(_fit)

            if _fit < self.p_bestFit[i]:
                self.p_best[i] = self.x[i]
                self.p_bestFit[i] = _fit
            if _fit < self.g_bestFit:
                self.g_best = self.x[i]
                self.g_bestFit = _fit
        ax1.scatter(x,y,color='red',marker='o',) #函数曲线上渲染种群点
            
    
    def pso(self,draw = 1):
        best_fit = [] #记录每轮迭代最佳适应度，用于绘图
        w_range = None
        fig=plt.figure(num=1,figsize=(4,4))
        ax1=fig.add_subplot(211)  #函数曲线图
        ax2=fig.add_subplot(212)  #收敛图
        if isinstance(self.w,tuple):
            w_range = self.w[1] - self.w[0]
            self.w = self.w[1]
        time_start = time.time() #记录迭代寻优开始时间
        for i in range(self.M):
            
            ax1.cla()
            x = np.linspace(0,50,500)
            y = x * np.sin(x) * np.cos(x*2) - 2*x*np.sin(3*x) + 3*x*np.sin(4*x)
            ax1.plot(x,y,color='blue')
            self.update(ax1)
            if w_range:
                self.w -= w_range / self.M
            print("\rIter：{:d}/{:d} fitness:{:.4f}".format(i,self.M,self.g_bestFit,end = '\n'))
            best_fit.append(self.g_bestFit.copy())
            ax2.plot(best_fit)   
            plt.pause(0.1)
        plt.pause(0)  #阻止迭代后图表消失的问题
        time_end = time.time() #记录迭代结束时间
        print(f'共花费 {time_end - time_start} 秒')
        
            

    

if __name__ == '__main__':
    low = [-50, -50, -50, -50, -50]
    up = [50, 50, 50, 50, 50]
    pso = PSO(30, 30, 300, low, up, -10, 10, w = 0.5)
    pso.pso()

